#!/usr/bin/env python3
#coding:utf-8

__author__ = 'xmxoxo<xmxoxo@qq.com>'


# 图像格式之间的转换：文件, bytes, BytesIO, PIL.Image, np.array, base64

from baselibs import *
import base64
import cv2
import numpy as np
from io import BytesIO
from PIL import Image
import win32clipboard

def get_base64(filename):
    ''' 将指定文件名的图像 转为base64编码
    '''
    try:
        with open(filename, 'rb') as f:
            base64_data = base64.b64encode(f.read())
            s = base64_data.decode()
            return s
    except Exception as e:
        return ''

def get_base64_stream(stream):
    base64_data = base64.b64encode(stream)
    s = base64_data.decode()
    return s

def array2base64(image_array):
    ''' np.ndarray ==>base64
    '''
    success, encoded_image = cv2.imencode(".jpg", image_array)
    if not success:
        return ''

    # 对数组的图片格式进行编码
    bytes = encoded_image.tostring()
    # print('bytes:', type(bytes), len(bytes))
    b64 = get_base64_stream(bytes)
    return b64

def base64toarray(image_b64):
    ''' base64 ==> np.ndarray
    '''
    # 解码
    bytes = base64.b64decode(image_b64)
    # print('bytes:', type(bytes), len(bytes))
    # 转 ndarray
    ret = cv2.imdecode(np.frombuffer(bytes, np.uint8), cv2.IMREAD_COLOR)
    image = cv2.cvtColor(ret, cv2.COLOR_BGR2RGB)
    return image

class IMGBase64():
    def __init__(self):
        self.filename = ''
        self.b64 = ''
        self.ndarray = None

    def set_filename(self, filename):
        self.filename = filename
        self.b64 = get_base64(filename)
        self.ndarray = None

    def set_array(self, image_array):
        self.ndarray = image_array
        self.b64 = array2base64(image_array)
        self.filename = ''

    def get_array(self):
        if self.ndarray is None:
            if self.b64 != '':
                image = base64toarray(self.b64)
                self.ndarray = image
        return self.ndarray

    def get_base64(self):
        return self.b64

    def save_file(self, filename):
        if self.b64 != '':
            with open(filename, 'wb') as f:
                bytes = base64.b64decode(self.b64)
                f.write(bytes)
            return 1
        else:
            return 0

    def paste(self):
        ''' 把图像粘贴到剪贴板, 需要 win32clipboard 包
        '''

        if self.b64 != '':
            # 解码base64数据
            bytes = base64.b64decode(self.b64)
            # 创建BytesIO对象
            image_io = BytesIO(bytes)
            # 打开图像
            image = Image.open(image_io)

            # 声明output字节对象
            output = BytesIO()
            # 用BMP (Bitmap) 格式存储
            # 这里是位图，然后用output字节对象来存储
            image.save(output, 'BMP')
            # BMP图片有14字节的header，需要额外去除
            data = output.getvalue()[14:]
            # 关闭
            output.close()

            # DIB: 设备无关位图(device-independent bitmap)，名如其意
            # BMP的图片有时也会以.DIB和.RLE作扩展名
            # 设置好剪贴板的数据格式，再传入对应格式的数据，才能正确向剪贴板写入数据
            send_msg_to_clip(win32clipboard.CF_DIB, data)
            print("图像已复制到剪贴板。")
            return 1
        else:
            return 0

def send_msg_to_clip(type_data, msg):
    """
    操作剪贴板分四步：
    1. 打开剪贴板：OpenClipboard()
    2. 清空剪贴板，新的数据才好写进去：EmptyClipboard()
    3. 往剪贴板写入数据：SetClipboardData()
    4. 关闭剪贴板：CloseClipboard()

    :param type_data: 数据的格式，unicode字符通常是传 win32con.CF_UNICODETEXT
    :param msg: 要写入剪贴板的数据
    """
    win32clipboard.OpenClipboard()
    win32clipboard.EmptyClipboard()
    if not msg is None:
        win32clipboard.SetClipboardData(type_data, msg)
    win32clipboard.CloseClipboard()

def test_IMGBase64():
    ''' 单元测试
    '''
    fname = 'image_cap_0.png'

    imgb64 = IMGBase64()
    imgb64.set_filename(fname)
    b64 = imgb64.get_base64()
    print('base64: %d bytes, string:%s' % (len(b64), b64[:50]))
    savetofile(b64, 'base64.txt')

    arr = imgb64.get_array()
    print('array:', arr.shape)

    nfname = 'tmp.png'
    #ret = imgb64.save_file(nfname)
    #print('保存文件:', ret)

    ret = imgb64.paste()

if __name__ == '__main__':
    pass
    test_IMGBase64()